/** \file    avi.cpp
  * \brief   AVI-streams facilities.
  * \date    2003
  * \author  Fedor Chelnokov, Pavel Agapov
  * \version $Header: /home/cscom1r4/cvsroot/cs/imglib/src/avi.cpp,v 1.5 2004/06/16 18:37:06 cscom1r4 Exp $
  */

#include "imglib/avi.h"

#pragma comment(lib, "Vfw32.lib")

namespace img {

namespace avi {

class initializer {
public:
	initializer()
		{ _ref_count = 0; }
	void init()
		{ if (++_ref_count == 1) AVIFileInit(); }
	void exit()
		{ if (--_ref_count == 0) AVIFileExit(); }
private:
	int _ref_count;
};

static initializer the_initializer;

user::user()
{
	the_initializer.init();
}

user::~user()
{
	the_initializer.exit();
}

const char * decode_icerr(DWORD e)
{
	if (e < ICERR_CUSTOM)
		return "custom codec error";
	switch (e) {
	case ICERR_OK:
		return "ok (no error)";
	case ICERR_DONTDRAW: 
		return "don't draw";
	case ICERR_NEWPALETTE:
		return "new palette";
	case ICERR_GOTOKEYFRAME:
		return "go to keyfame";
	case ICERR_STOPDRAWING:
		return "stop drawing";
	case ICERR_UNSUPPORTED:
		return "unsupported";
	case ICERR_BADFORMAT:
		return "bad format";
	case ICERR_MEMORY:
		return "memory (out of)";
	case ICERR_INTERNAL:
		return "internal error";
	case ICERR_BADFLAGS:
		return "bad flags";
	case ICERR_BADPARAM:
		return "bad param";
	case ICERR_BADSIZE:
		return "bad size";
	case ICERR_BADHANDLE:
		return "bad handle";
	case ICERR_CANTUPDATE:
		return "can't update";
	case ICERR_ABORT:
		return "abort";
	case ICERR_ERROR:
		return "error";
	case ICERR_BADBITDEPTH:
		return "bad bit depth";
	case ICERR_BADIMAGESIZE:
		return "bad image size";
	default:
		return "unknown error";
	};
}

///////////
// coder //
///////////

codec_id & codec_id::operator = (std::string const & name)
{
	if (name.length() >= 4)
		std::copy(name.begin(), name.begin()+4, code);
	else {
		std::copy(name.begin(), name.end(), code);
		std::fill(code + name.length(), code + 4, char(0));
	}
	return * this;
}

coder::coder()
{
	_pci = 0;
}

coder::coder(const coder & rhs)
{
	_pci = 0;
	* this = rhs;
}

coder::coder(const coder_info & ci)
{
	_pci = 0;
	open(ci);
}

coder::~coder()
{
	close();
}

coder & coder::operator = (const coder & rhs)
{
	if (rhs.is_open()) {
		open(*rhs.info());
		img::avi::state st;
		rhs.state(st);
		set_state(st);
	}
	else 
		close();
	return * this;
}

void coder::open(const coder_info & ci)
{
	if (is_open())
		close();
	_hic = ICOpen(ICTYPE_VIDEO, ci.id.fccHandler, ICMODE_COMPRESS);
	if (_hic == 0)
		throw exception("error opening video compressor");

	_pci = &ci;
}

void coder::open(coder_info & ci, DWORD fccHandler, BITMAPINFOHEADER * in_format)
{
	if (is_open())
		close();
	_hic = ICOpen(ICTYPE_VIDEO, fccHandler, ICMODE_COMPRESS);
	if (_hic == 0)
		throw exception("error opening video compressor");

	ci.interrogate(_hic, in_format);
	_pci = &ci;
}

bool coder::choose(coder_info & ci, HWND hwnd, BITMAPINFOHEADER * in_format, PAVISTREAM preview)
{
	if (is_open())
		close();

	UINT flags = ICMF_CHOOSE_DATARATE | ICMF_CHOOSE_KEYFRAME;
	if (!in_format)
		flags |= ICMF_CHOOSE_ALLCOMPRESSORS;
	if (preview)
		flags |= ICMF_CHOOSE_PREVIEW;
	COMPVARS c;
	memset(&c, 0, sizeof(c));
	c.cbSize = sizeof(c);
	if (FALSE == ICCompressorChoose(hwnd, flags, in_format, preview, &c, 0)) {
		ICCompressorFree(&c);
		return false;
	}

	ci.interrogate(c.hic, in_format);
	_pci = &ci;
	_hic = c.hic;
	avi::state s;
	state(s);
	_hic = 0;
	ICCompressorFree(&c);

	open(ci);
	set_state(s);

	return true;
}

void coder::close()
{
	if (is_open()) {
		ICClose(_hic);
		_pci = 0;
	}
}

void coder::set_state(const avi::state & s)
{
	assert(is_open());
	//DIVX and XVID return zero here -- very strage :-(
	//if (!ICSetState(_hic, &*s.begin(), s.size()) && s.size() != 0)
	//	throw exception("error setting state of video compressor");
	ICSetState(_hic, &*s.begin(), s.size());
}

void coder::state(avi::state & s) const
{
	assert(is_open());
	DWORD size = ICGetStateSize(_hic);
	s.resize(size);
	if (!size)
		return;
	//most coders (including MS ones) do not follow documentation: 
	//ICGetState returns size instead of ICERR_OK 
	DWORD res = ICGetState(_hic, &*s.begin(), size);
	if (size != res && ICERR_OK != res) 
		throw exception("error getting state of video compressor");
}

void coder::configure(HWND hwnd)
{
	assert(is_open());
	assert(_pci->configurable);
	ICConfigure(_hic, hwnd);
}

void coder::about(HWND hwnd)
{
	assert(is_open());
	assert(_pci->has_about_dialog);
	ICAbout(_hic, hwnd);
}

/////////////////
// enum_coders //
/////////////////

void format::fill(BITMAPINFO & info) const
{
	memset(&info, 0, sizeof(info));
	info.bmiHeader.biSize = sizeof(BITMAPINFOHEADER);
	info.bmiHeader.biWidth = dim.sx;
	info.bmiHeader.biHeight = dim.sy; //most codecs does not support negative values here,
								     //so bottom_up order of lines is requirement
	info.bmiHeader.biBitCount = alpha ? 32 : 24;
	info.bmiHeader.biPlanes = 1;
}

void coder_info::interrogate(HIC hic, BITMAPINFOHEADER * in_format)
{
	BITMAPINFO my_info;
	if (!in_format) {
		format(size_extent(720,576), fraction(25,1)).fill(my_info);
		in_format = &my_info.bmiHeader;
	}

	int bc = in_format->biBitCount;
	in_format->biBitCount = 32;
	alpha = ICCompressQuery(hic, in_format, NULL) == ICERR_OK;
	in_format->biBitCount = 24;
	noalpha = ICCompressQuery(hic, in_format, NULL) == ICERR_OK;
	in_format->biBitCount = bc;

	ICINFO icinfo;
	ICGetInfo(hic, &icinfo, sizeof(icinfo)); //it's strange but ICInfo(...) does not fill icinfo!
	id.fccHandler = icinfo.fccHandler;
	name.assign(icinfo.szName, icinfo.szName+16);
	name = name.substr(0, name.find('\0'));
	description.assign(icinfo.szDescription, icinfo.szDescription+128);
	description = description.substr(0, description.find('\0'));
	quality = (icinfo.dwFlags & VIDCF_QUALITY) != 0;
	temporal = (icinfo.dwFlags & VIDCF_TEMPORAL) != 0;
	crunch = (icinfo.dwFlags & VIDCF_CRUNCH) != 0;
	has_about_dialog = ICQueryAbout(hic);
	configurable = ICQueryConfigure(hic);
}

void enum_coders(coder_infos & a, size_extent e)
{
	format f(e, fraction(25,1));
	BITMAPINFO info;
	f.fill(info);

	ICINFO icinfo;
	for (int i=0; ICInfo(ICTYPE_VIDEO, i, &icinfo); i++) 
	{ 
#ifdef _DEBUG
		//DivX has some protection against running under debuger
		if (icinfo.fccHandler == 2021026148)
			continue;
#endif
		HIC hic = ICOpen(ICTYPE_VIDEO, icinfo.fccHandler, ICMODE_QUERY); 
		if (hic) 
		{ 
			coder_info ci;
			ci.interrogate(hic, &info.bmiHeader);
			// Skip this compressor if it can't handle the format. 
			if (ci.alpha || ci.noalpha)
				a.push_back(ci);
 
			ICClose(hic); 
		} 
	} 
}

////////////
// writer //
////////////

writer::writer(coder & c, const std::string & filename, const format & f) :
	_coder(c)
{
	init();
	try {
		open(filename, f);
	}
	catch (...) {
		close();
		throw;
	}
}

writer::writer(coder & c) :
	_coder(c)
{
	init();
}

void writer::init()
{
	_pfile = 0;
	_ps = 0;
	_lpbiOut = 0;
	_outBuffer = 0;
}

writer::~writer()
{
	close();
}

void writer::open(const std::string & filename, const format & f)
{
	assert(_coder.is_open());
	if (_ps || _pfile || _lpbiOut || _outBuffer)
		close();

	//changes alpha-channel if it is not supported by the coder
	_fmt = f;
	if (f.alpha && !_coder.info()->alpha)
		_fmt.alpha = false;
	else if (!f.alpha && !_coder.info()->noalpha)
		_fmt.alpha = true;

	_fmt.fill(_in);

	if (_fmt.alpha) 
		_scansize = _fmt.dim.sx * 4;
	else
		_scansize = (_fmt.dim.sx * 24 + 31) / 32 * 4;

	DWORD dwFormatSize = ICCompressGetFormatSize(hic(), &_in);
	_lpbiOut = (LPBITMAPINFOHEADER) new char[dwFormatSize]; 
	ICCompressGetFormat(hic(), &_in, _lpbiOut); 

	// Find the worst-case buffer size. 
	DWORD dwCompressBufferSize = ICCompressGetSize(hic(), &_in, _lpbiOut); 
	// Allocate a buffer and get _outBuffer to point to it. 
	_outBuffer = new char[dwCompressBufferSize]; 

	if (AVIERR_OK != AVIFileOpen(
		&_pfile,				// returned file pointer
		filename.c_str(),		// file name
		OF_WRITE | OF_CREATE | OF_SHARE_EXCLUSIVE, // mode to open file with
		NULL))					// use handler determined
		throw exception("AVI file open error");

	AVISTREAMINFO header;

	memset(&header, 0, sizeof(header));
	header.fccType					= streamtypeVIDEO;// stream type
	header.fccHandler				= _coder.info()->id.fccHandler;
	header.dwScale					= _fmt.fps.den;
	header.dwRate					= _fmt.fps.num;
	header.dwSuggestedBufferSize	= dwCompressBufferSize;
	header.rcFrame.right			= _fmt.dim.sx;
	header.rcFrame.bottom			= _fmt.dim.sy;

	// And create the stream;
	if (AVIERR_OK != AVIFileCreateStream(
		_pfile,				// file pointer
		&_ps,				// returned stream pointer
		&header))			// stream header
		throw exception("AVI stream in file create error");

	if (AVIERR_OK != AVIStreamSetFormat(_ps, 0,	_lpbiOut, dwFormatSize))
		throw exception("AVI stream set format error");

	if (ICERR_OK != ICCompressBegin(hic(), &_in, _lpbiOut))
		throw exception("video compression begin error");

	_frames_out = 0;
}

void writer::close()
{
	ICCompressEnd(hic());  // terminate compression 

	if (_ps)
		AVIStreamClose(_ps), _ps = 0;

	if (_pfile)
		AVIFileClose(_pfile), _pfile = 0;

	if (_lpbiOut) {
		delete[] (char*)_lpbiOut;
		_lpbiOut = 0;
	}

	if (_outBuffer) {
		delete[] _outBuffer;
		_outBuffer = 0;
	}

	_prev32.clear();
	_prev24.clear();
}

writer & writer::operator << (const byte_image & img)
{
	assert(is_open());
	assert(img.dim() == _fmt.dim);
	if (img.scansize() == _scansize && img.order() == bottom_up && !_coder.info()->temporal)
		write_frame(img.data());
	else {
		if (_fmt.alpha) {
			byte_image i32;
			i32.order(bottom_up);
			i32.resize(img.dim());
			i32 = img;
			assert(i32.scansize() == _scansize && i32.order() == bottom_up);
			write_frame(i32.data(), _prev32.data());
			if (_coder.info()->temporal)
				_prev32.swap(i32);
		}
		else {
			small_byte_image i24;
			i24.order(bottom_up);
			i24.resize(img.dim());
			i24 = img;
			assert(i24.scansize() == _scansize && i24.order() == bottom_up);
			write_frame(i24.data(), _prev24.data());
			if (_coder.info()->temporal)
				_prev24.swap(i24);
		}
	}
	return * this;
}

writer & writer::operator << (const small_byte_image & img)
{
	assert(is_open());
	assert(img.dim() == _fmt.dim);
	assert(!_fmt.alpha);
	if (img.scansize() == _scansize && img.order() == bottom_up && !_coder.info()->temporal)
		write_frame(img.data());
	else {
		small_byte_image i24;
		i24.order(bottom_up);
		i24.resize(img.dim());
		i24 = img;
		assert(i24.scansize() == _scansize && i24.order() == bottom_up);
		write_frame(i24.data(), _prev24.data());
		if (_coder.info()->temporal)
			_prev24.swap(i24);
	}
	return * this;
}

void writer::write_frame(const void * data, const void * prev_data)
{
	DWORD dwFlags;
	DWORD ckid;
	if (ICERR_OK != ICCompress(hic(), 0, 
		_lpbiOut, _outBuffer, 
		&_in.bmiHeader, (LPVOID)data, 
		&ckid, &dwFlags,
		_frames_out, 0, 0, 
		prev_data ? &_in.bmiHeader : 0, (LPVOID)prev_data))
		throw exception("video compression error");

	//writing data for frame #150 I received AVIERR_UNSUPPORTED error from AVIStreamWriteData()
	//moreover DIVX (and some other coders) does(do) not fill ckid in ICCompress,
	//while Cinepack places one and the same constant in ckid

	//HRESULT h;
	//if ((h = AVIStreamWriteData(_ps, ckid, 0, 0)))
	//	throw exception("AVI stream data write error");

	if (AVIERR_OK != AVIStreamWrite(
		_ps,				// stream pointer
		_frames_out++,		// time of this frame
		1,					// number to write
		(LPBYTE) _outBuffer,
		_lpbiOut->biSizeImage,
		dwFlags,			//AVIIF_KEYFRAME, flags....
		NULL,
		NULL))
		throw exception("AVI stream write error");
}

} //namespace avi

} //namespace img
